from mmengine.config import read_base
with read_base():
    from .mgam import *

from mmengine.model.weight_init import PretrainedInit
from itkit.mm.mmseg_Dev3D import DiceLoss_3D
from itkit.mm.mgam_models import mgam_Seg3D_Lite
from itkit.models.SegFormer3D import SegFormer3D

model = dict(
    type=mgam_Seg3D_Lite,
    binary_segment_threshold=None,
    num_classes=num_classes,
    inference_PatchSize=size,
    inference_PatchStride=[s//2 for s in size],
    inference_PatchAccumulateDevice='cpu',
    backbone=dict(
        type=SegFormer3D,
        in_channels=in_channels, # pyright: ignore
        num_classes=num_classes,
        sr_ratios=[4, 2, 1, 1],
        embed_dims=[32, 64, 128, 256],
        num_heads=[4, 8, 16, 16],
        depths=[2, 2, 2, 2],
        mlp_ratios=[2, 2, 2, 2],
        decoder_head_embedding_dim=64,
    ),
    criterion=dict(
        type=DiceLoss_3D,
        to_onehot_y=True,
        sigmoid=True,
        softmax=False,
        batch=False,
    )
)
